import {Annotation, MemorySaver, START, StateGraph} from "@langchain/langgraph";

const SubgraphStateAnnotation = Annotation.Root({
    foo: Annotation<string>,
    bar: Annotation<string>
})

const subgraphNode1 = async (state: typeof SubgraphStateAnnotation.State) => {
    return {
        bar: "bar"
    }
}

const subgraphNode2 = async (state: typeof SubgraphStateAnnotation.State)=> {
    return {
        foo: state.foo + state.bar
    }
}

const subgraph = new StateGraph(SubgraphStateAnnotation)
.addNode('subgraphNode1', subgraphNode1)
.addNode('subgraphNode2', subgraphNode2)
.addEdge(START, 'subgraphNode1')
.addEdge('subgraphNode1', 'subgraphNode2')
.compile();

// parent graph
const StateAnnotation =  Annotation.Root({
    foo: Annotation<string>
});

const node1 = async (state: typeof StateAnnotation.State)=> {
    return {
        foo: "hi! " + state.foo
    }
}

const  builder = new StateGraph(StateAnnotation)
.addNode('node1', node1)
.addNode('node2', subgraph)
.addEdge(START, 'node1')
.addEdge('node1', 'node2')

const checkpointer = new MemorySaver();

const graph = builder.compile({
    checkpointer: checkpointer
})

const config = {
    configurable: {
        thread_id: '1'
    }
}

const stream = await graph.stream({
    foo: "foo"
}, {
    ...config,
    subgraphs: true
})

for await (const [_source, chunk] of stream) {
    console.log(chunk)
}

let stateWithSubgraph;

const graphHistories = await graph.getStateHistory(config);

for await (const state of graphHistories) {
    if (state.next[0] === "node2") {
        stateWithSubgraph = state;
        break;
    }
}
const subgraphConfig = stateWithSubgraph.tasks[0].state;
console.log(subgraphConfig)
console.log((await graph.getState(subgraphConfig)).values)